package com.zhuhjay.framework.helper;

import cn.hutool.core.util.StrUtil;
import com.zhuhjay.common.base.enums.LimiterEnums;
import com.zhuhjay.common.base.exception.ApiLimiterException;
import com.zhuhjay.common.base.model.domain.LimiterDetail;
import com.zhuhjay.common.core.helper.ApiLimiterHelper;
import com.zhuhjay.common.redis.util.RedisUtils;
import com.zhuhjay.service.system.service.SystemBlackService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.List;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import static com.zhuhjay.common.base.constant.ApiLimiterConstants.DEFAULT_MESSAGE;
import static com.zhuhjay.common.base.constant.ApiLimiterConstants.DEFAULT_NOTICE_COUNT;
import static com.zhuhjay.common.base.constant.CacheConstant.*;
import static com.zhuhjay.common.redis.context.RedisScriptContext.API_BLACK_LIMITER_SCRIPT;
import static com.zhuhjay.common.redis.context.RedisScriptContext.API_LIMITER_SCRIPT;

/**
 * @author ZhuHJay
 * @date 2023/7/27 22:57
 */
@Slf4j
@Component
public class ApiLimiterHelperImpl implements ApiLimiterHelper {

    @Resource
    private SystemBlackService blackService;

    @Override
    public void limiter(String limitKeySuf, LimiterDetail detail, String reason) {
        String target = detail.getTarget();
        String limitKey = StrUtil.format("{}:{}", LIMITER_PREFIX + target, limitKeySuf);
        String noticeKey = LIMITER_NOTICE_PREFIX + target;
        /* 使用阻塞分布式锁进行限流, 已淘汰
        RedissonUtils.withInfiniteWaitLock(limitKey, () -> {
            checkBlack(detail.getStrategy(), detail.getTarget());
            limiter(limitKey, detail.getLimit(), detail.getExpire(), detail.getTarget(), detail.getMessage());
        });*/
        // 相比于上面的方式, 使用 lua 脚本进行限流, 提高了 10 倍以上的 QPS
        Long count = RedisUtils.execute(
                API_BLACK_LIMITER_SCRIPT,
                List.of(limitKey, noticeKey, BLACK_KEY),
                String.valueOf(detail.getLimit()),
                String.valueOf(detail.getExpire()),
                String.valueOf(DEFAULT_NOTICE_COUNT),
                // 限流通知过期时间(单位:s)
                String.valueOf(3600),
                target
        );
        // 已经进入黑名单
        if (count.intValue() == -1) {
            throw new ApiLimiterException("您已被限制访问, 请联系管理员!");
        }
        // 超出限制提醒次数, 加入黑名单
        if (count.intValue() == detail.getLimit() + DEFAULT_NOTICE_COUNT) {
            blackService.addBlack4Limiter(detail.getStrategy(), target, reason);
            throw new ApiLimiterException("您已被限制访问, 请联系管理员!");
        }
        // 限流提醒
        String message = detail.getMessage();
        if (count.intValue() == 0) {
            throw new ApiLimiterException(StrUtil.isBlank(message) ? DEFAULT_MESSAGE : message);
        }
    }

    /**
     * 黑名单检查(是否在黑名单中, 是否超出限制提醒次数)
     * 黑名单：限制对应类型的访问
     * eg: 限制IP访问, 则将IP加入访问黑名单, 限制需要IP的访问 @ApiLimiter(strategy = IP)
     * 限制用户访问, 则将用户加入访问黑名单, 限制需要用户的访问 @ApiLimiter(strategy = USER)
     **/
    @Deprecated
    private void checkBlack(LimiterEnums type, String target) {
        // 是否存在于黑名单中
        Set<String> strings = RedisUtils.getMembers4Set(BLACK_KEY);
        if (strings.contains(target)) {
            throw new ApiLimiterException("您已被限制访问!");
        }
        // 是否超出限制提醒次数
        String cache = RedisUtils.getCache(LIMITER_NOTICE_PREFIX + target);
        int num = Integer.parseInt(StrUtil.isBlank(cache) ? "0" : cache);
        if (num > DEFAULT_NOTICE_COUNT) {
            blackService.addBlack4Limiter(type, target, "限流器自动限流");
            // 限制访问后清除提醒次数
            RedisUtils.deleteCache(LIMITER_NOTICE_PREFIX + target);
            throw new ApiLimiterException("访问过于频繁, 已被限制访问!");
        }
    }

    /**
     * 使用 redis 进行限流
     **/
    @Deprecated
    private void limiter(String key, int limit, int expire, String target, String message) {
        // 执行限流脚本并返回访问次数
        Long count = RedisUtils.execute(
                API_LIMITER_SCRIPT,
                List.of(key),
                String.valueOf(limit), String.valueOf(expire)
        );
        // 当返回值为 0 的时候则进行限流提醒
        if (count != null && count.intValue() == 0) {
            // 计数限流提醒次数
            Long num = RedisUtils.increment(LIMITER_NOTICE_PREFIX + target);
            // 如果是第一次则设置过期时间
            if (num != null && num.intValue() == 1) {
                RedisUtils.setKeyExpire(LIMITER_NOTICE_PREFIX + target, 1, TimeUnit.HOURS);
            }
            throw new ApiLimiterException(StrUtil.isBlank(message) ? DEFAULT_MESSAGE : message);
        }
    }
}
